/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <iostream>
#include <yhamg.h>

using namespace YHAMG;

int main(int argc, char* argv[])
{
	MPI_Init(&argc, &argv);

	int solver_type = 0;
	int pre_type = 0;
	char* in_file = 0;
	char* rhs_file = 0;
	char* out_file = 0;

	int arg_index = 1;

	while (arg_index < argc)
	{
		if (strcmp(argv[arg_index], "-in") == 0)
			in_file = argv[++arg_index];
		else if (strcmp(argv[arg_index], "-rhs") == 0)
			rhs_file = argv[++arg_index];
		else if (strcmp(argv[arg_index], "-out") == 0)
			out_file = argv[++arg_index];
		else if (strcmp(argv[arg_index], "-pre") == 0)
			pre_type = atoi(argv[++arg_index]);
		else if (strcmp(argv[arg_index], "-solver") == 0)
			solver_type = atoi(argv[++arg_index]);
		++arg_index;
	}

	int    max_levels = 20;
	int    coarse_size = 8;
	double strength_threshold = 0.25;
	int    agg_num_levels = 1;
	int    coarsen_type = 0;
	int    interp_min_elements = 4;
	int    interp_max_elements = 4;
	double truncation_factor = 0.1;
	double sparsification_threshold = 0.01;
	int    cycle_type = 0;
	int    pre_sweeps = 1;
	int    post_sweeps = 1;
	int    coarse_sweeps = 2;
	int    relaxation_type = 3;
	double relaxation_factor = 1.0;

	int    sor_type = 2;
	double sor_factor = 1.0;

	int    gmres_restart = 10;
	int    max_iters = 500;
	double tolerance = 1.0e-08;

	bool   print_stats = 1;

	if (!in_file)
		return -1;

	int MPI_RANK;
	MPI_Comm_rank(MPI_COMM_WORLD, &MPI_RANK);

	if (MPI_RANK == 0)
	{
		MPI_Comm comm = MPI_COMM_SELF;

		ParOpBase* A;
		ParOpBase* P;
		ParVec x(comm);
		ParVec b(comm);

		int iter;
		double relres;
		double setup_time;
		double solve_time;

		char line[1024];
		char banner[64];
		char mtx[64]; 
		char crd[64];
		char data_type[64];
		char storage_scheme[64];
		
		FILE* fin = fopen(in_file, "r");

		fgets(line, 1024, fin);
		sscanf(line, "%s %s %s %s %s", &banner, &mtx, &crd, &data_type, &storage_scheme);

		char* p;
		p = data_type;
		*p = toupper(*p);
		while (*(++p) != '\0')
			*p = tolower(*p);

		p = storage_scheme;
		*p = toupper(*p);
		while (*(++p) != '\0')
			*p = tolower(*p);

		do
			fgets(line, 1024, fin);
		while (line[0] == '%');

		if (strcmp(data_type, "Real") == 0)
		{
			if (strcmp(storage_scheme, "General") == 0)
			{
				int N, M, Nz;
				sscanf(line, "%d %d %d", &N, &M, &Nz);

				int* I = new int[Nz];
				int* J = new int[Nz];
				double* V = new double[Nz];

				for (int j = 0; j < Nz; ++j)
				{
					fscanf(fin, "%d %d %lg\n", &I[j], &J[j], &V[j]);
					--I[j];
					--J[j];
				}

				fclose(fin);
				
				A = new ParCSR;
				((ParCSR*)A)->comm = comm;
				((ParCSR*)A)->LocLoc = COO(N, M, Nz, I, J, V, 0);
				((ParCSR*)A)->LocExt.size[0] = N;
				((ParCSR*)A)->LocExt.rowptr = new int[N + 1];
				for (int i = 0; i <= N; ++i)
					((ParCSR*)A)->LocExt.rowptr[i] = 0;
				((ParCSR*)A)->Ext_ptr = new int[1];
				((ParCSR*)A)->Ext_ptr[0] = 0;
			}
			else if (strcmp(storage_scheme, "Symmetric") == 0)
			{
				int N, M, Nz, _Nz;
				sscanf(line, "%d %d %d", &N, &M, &_Nz);

				int* I = new int[_Nz * 2];
				int* J = new int[_Nz * 2];
				double* V = new double[_Nz * 2];
				
				Nz = 0;
				for (int j = 0; j < _Nz; ++j)
				{
					fscanf(fin, "%d %d %lg\n", &I[Nz], &J[Nz], &V[Nz]);
					--I[Nz];
					--J[Nz];
					if (I[Nz] != J[Nz])
					{
						I[Nz + 1] = J[Nz];
						J[Nz + 1] = I[Nz];
						V[Nz + 1] = V[Nz];
						++Nz;
					}
					++Nz;
				}

				fclose(fin);
				A = new ParCSR;
				((ParCSR*)A)->comm = comm;
				((ParCSR*)A)->LocLoc = COO(N, M, Nz, I, J, V, 0);
				((ParCSR*)A)->LocExt.size[0] = N;
				((ParCSR*)A)->LocExt.rowptr = new int[N + 1];
				for (int i = 0; i <= N; ++i)
					((ParCSR*)A)->LocExt.rowptr[i] = 0;
				((ParCSR*)A)->Ext_ptr = new int[1];
				((ParCSR*)A)->Ext_ptr[0] = 0;
			}

			std::cout << "\nTest Problem:\n";
			std::cout << "----------------------------------------------------------------\n";
			std::cout << "Input File: \"" << in_file << "\"\n";
			std::cout << "Type & Structure: " <<  data_type << " " << storage_scheme << "\n";
			std::cout << "Rows: " << ((ParCSR*)A)->LocLoc.size[0] << ", Nz: " << ((ParCSR*)A)->LocLoc.rowptr[((ParCSR*)A)->LocLoc.size[0]] << "\n";
			
			((ParCSR*)A)->SetupHalo();

			if (pre_type == 0)
			{
				P = new ParCSRPrecondAMG();

				((ParCSRPrecondAMG*)P)->MaxLevels = max_levels;
				((ParCSRPrecondAMG*)P)->CoarseSize = coarse_size;
				((ParCSRPrecondAMG*)P)->StrengthThreshold = strength_threshold;
				((ParCSRPrecondAMG*)P)->AggNumLevels = agg_num_levels;
				((ParCSRPrecondAMG*)P)->CoarsenType = coarsen_type;
				((ParCSRPrecondAMG*)P)->InterpMinElements = interp_min_elements;
				((ParCSRPrecondAMG*)P)->InterpMaxElements = interp_max_elements;
				((ParCSRPrecondAMG*)P)->TruncationFactor = truncation_factor;
				((ParCSRPrecondAMG*)P)->SparsificationThreshold = sparsification_threshold;
				((ParCSRPrecondAMG*)P)->CycleType = cycle_type;
				((ParCSRPrecondAMG*)P)->PreSweeps = pre_sweeps;
				((ParCSRPrecondAMG*)P)->PostSweeps = post_sweeps;
				((ParCSRPrecondAMG*)P)->CoarseSweeps = coarse_sweeps;
				((ParCSRPrecondAMG*)P)->RelaxationType = relaxation_type;
				((ParCSRPrecondAMG*)P)->RelaxationFactor = relaxation_factor;
				((ParCSRPrecondAMG*)P)->PrintStats = print_stats;
			}
			else if (pre_type == 1)
			{
				P = new ParCSRPrecondILUAMG;

				((ParCSRPrecondILUAMG*)P)->MaxLevels = max_levels;
				((ParCSRPrecondILUAMG*)P)->CoarseSize = coarse_size;
				((ParCSRPrecondILUAMG*)P)->StrengthThreshold = strength_threshold;
				((ParCSRPrecondILUAMG*)P)->AggNumLevels = agg_num_levels;
				((ParCSRPrecondILUAMG*)P)->CoarsenType = coarsen_type;
				((ParCSRPrecondILUAMG*)P)->InterpMinElements = interp_min_elements;
				((ParCSRPrecondILUAMG*)P)->InterpMaxElements = interp_max_elements;
				((ParCSRPrecondILUAMG*)P)->TruncationFactor = truncation_factor;
				((ParCSRPrecondILUAMG*)P)->SparsificationThreshold = sparsification_threshold;
				((ParCSRPrecondILUAMG*)P)->CycleType = cycle_type;
				((ParCSRPrecondILUAMG*)P)->PreSweeps = pre_sweeps;
				((ParCSRPrecondILUAMG*)P)->PostSweeps = post_sweeps;
				((ParCSRPrecondILUAMG*)P)->CoarseSweeps = coarse_sweeps;
				((ParCSRPrecondILUAMG*)P)->PrintStats = print_stats;
			}
			else if (pre_type == 2)
				P = new ParCSRPrecondJacobi;
			else if (pre_type == 3)
				P = new ParCSRPrecondSOR(sor_type, sor_factor);
			else if (pre_type == 4)
				P = new ParCSRPrecondILU;

			setup_time = -MPI_Wtime();

			if (pre_type == 0)
				std::cout << "\nAMG Setup Phase:\n";
			else if (pre_type == 1)
				std::cout << "\nILUAMG Setup Phase:\n";
			else if (pre_type == 2)
				std::cout << "\nJacobi Setup Phase:\n";
			else if (pre_type == 3)
				std::cout << "\nSOR Setup Phase:\n";
			else if (pre_type == 4)
				std::cout << "\nILU Setup Phase:\n";
			
			std::cout << "----------------------------------------------------------------\n";

			((ParCSRPrecondBase*)P)->Setup(*(ParCSR*)A);

			setup_time += MPI_Wtime();

			std::cout << "Setup Time: " << setup_time << "\n";
			
			if (!rhs_file)
			{
				b.Resize(A->OutSize());
				b.Fill(1.0);
			}
			else
			{
				FILE* frhs = fopen(rhs_file, "r");

				do
					fgets(line, 1024, frhs);
				while (line[0] == '%');

				int N, M;
				sscanf(line, "%d %d", &N, &M);

				double* bv = new double[N];
				for (int i = 0; i < N; ++i)
					fscanf(frhs, "%lg\n", &bv[i]);
				
				fclose(frhs);
				
				b = ParVec(comm, N, bv, 0);
			}
		}
		else if (strcmp(data_type, "Complex") == 0)
		{
			if (strcmp(storage_scheme, "General") == 0)
			{
				int N, M, Nz;
				sscanf(line, "%d %d %d", &N, &M, &Nz);

				int* I = new int[Nz];
				int* J = new int[Nz];
				double* V = new double[Nz * 2];

				for (int j = 0; j < Nz; ++j)
				{
					fscanf(fin, "%d %d %lg %lg\n", &I[j], &J[j], &V[2 * j], &V[2 * j + 1]);
					--I[j];
					--J[j];
				}
				
				fclose(fin);
				
				int* Ap = new int[N + 1];
				int* Ai = new int[Nz];
				double* Av = new double[4 * Nz];

				for (int i = 0; i <= N; ++i)
					Ap[i] = 0;

				for (int j = 0; j < Nz; ++j)
					++Ap[I[j]];

				for (int i = 0; i < N; ++i)
					Ap[i + 1] += Ap[i];

				for (int j = Nz - 1; j >= 0; --j)
				{
					int r = --Ap[I[j]];
					Ai[r] = J[j];
					Av[r * 4] = V[2 * j];
					Av[r * 4 + 1] = -V[2 * j + 1];
					Av[r * 4 + 2] = V[2 * j + 1];
					Av[r * 4 + 3] = V[2 * j];
				}
				
				A = new ParBSR;
				((ParBSR*)A)->comm = comm;
				((ParBSR*)A)->LocLoc.size[0] = N;
				((ParBSR*)A)->LocLoc.size[1] = M;
				((ParBSR*)A)->LocLoc.block_size = 2;
				((ParBSR*)A)->LocLoc.rowptr = Ap;
				((ParBSR*)A)->LocLoc.colind = Ai;
				((ParBSR*)A)->LocLoc.values = Av;
				((ParBSR*)A)->LocExt.size[0] = N;
				((ParBSR*)A)->LocExt.block_size = 2;
				((ParBSR*)A)->LocExt.rowptr = new int[N + 1];
				for (int i = 0; i <= N; ++i)
					((ParBSR*)A)->LocExt.rowptr[i] = 0;
				((ParBSR*)A)->Ext_ptr = new int[1];
				((ParBSR*)A)->Ext_ptr[0] = 0;
			}
			else if (strcmp(storage_scheme, "Hermitian") == 0)
			{
				int N, M, Nz, _Nz;
				sscanf(line, "%d %d %d", &N, &M, &_Nz);

				int* I = new int[_Nz * 2];
				int* J = new int[_Nz * 2];
				double* V = new double[_Nz * 4];

				Nz = 0;
				for (int j = 0; j < _Nz; ++j)
				{
					fscanf(fin, "%d %d %lg %lg\n", &I[Nz], &J[Nz], &V[2 * Nz], &V[2 * Nz + 1]);
					--I[Nz];
					--J[Nz];
					if (I[Nz] != J[Nz])
					{
						I[Nz + 1] = J[Nz];
						J[Nz + 1] = I[Nz];
						V[2 * (Nz + 1)] = V[2 * Nz];
						V[2 * (Nz + 1) + 1] = -V[2 * Nz + 1];
						++Nz;
					}
					++Nz;
				}
				
				fclose(fin);
				
				int* Ap = new int[N + 1];
				int* Ai = new int[Nz];
				double* Av = new double[4 * Nz];

				for (int i = 0; i <= N; ++i)
					Ap[i] = 0;

				for (int j = 0; j < Nz; ++j)
					++Ap[I[j]];

				for (int i = 0; i < N; ++i)
					Ap[i + 1] += Ap[i];

				for (int j = Nz - 1; j >= 0; --j)
				{
					int r = --Ap[I[j]];
					Ai[r] = J[j];
					Av[r * 4] = V[2 * j];
					Av[r * 4 + 1] = V[2 * j + 1];
					Av[r * 4 + 2] = -V[2 * j + 1];
					Av[r * 4 + 3] = V[2 * j];
				}
				
				A = new ParBSR;
				((ParBSR*)A)->comm = comm;
				((ParBSR*)A)->LocLoc.size[0] = N;
				((ParBSR*)A)->LocLoc.size[1] = M;
				((ParBSR*)A)->LocLoc.block_size = 2;
				((ParBSR*)A)->LocLoc.rowptr = Ap;
				((ParBSR*)A)->LocLoc.colind = Ai;
				((ParBSR*)A)->LocLoc.values = Av;
				((ParBSR*)A)->LocExt.size[0] = N;
				((ParBSR*)A)->LocExt.block_size = 2;
				((ParBSR*)A)->LocExt.rowptr = new int[N + 1];
				for (int i = 0; i <= N; ++i)
					((ParBSR*)A)->LocExt.rowptr[i] = 0;
				((ParBSR*)A)->Ext_ptr = new int[1];
				((ParBSR*)A)->Ext_ptr[0] = 0;
			}

			std::cout << "\nTest Problem:\n";
			std::cout << "----------------------------------------------------------------\n";
			std::cout << "Input File: \"" << in_file << "\"\n";
			std::cout << "Type & Structure: " <<  data_type << " " << storage_scheme << "\n";
			std::cout << "Rows: " << ((ParBSR*)A)->LocLoc.size[0] << ", Nz: " << ((ParBSR*)A)->LocLoc.rowptr[((ParBSR*)A)->LocLoc.size[0]] << "\n";

			((ParBSR*)A)->SetupHalo();

			if (pre_type == 0)
			{
				P = new ParBSRPrecondAMG();

				((ParBSRPrecondAMG*)P)->MaxLevels = max_levels;
				((ParBSRPrecondAMG*)P)->CoarseSize = coarse_size;
				((ParBSRPrecondAMG*)P)->StrengthThreshold = strength_threshold;
				((ParBSRPrecondAMG*)P)->AggNumLevels = agg_num_levels;
				((ParBSRPrecondAMG*)P)->CoarsenType = coarsen_type;
				((ParBSRPrecondAMG*)P)->InterpMinElements = interp_min_elements;
				((ParBSRPrecondAMG*)P)->InterpMaxElements = interp_max_elements;
				((ParBSRPrecondAMG*)P)->TruncationFactor = truncation_factor;
				((ParBSRPrecondAMG*)P)->SparsificationThreshold = sparsification_threshold;
				((ParBSRPrecondAMG*)P)->CycleType = cycle_type;
				((ParBSRPrecondAMG*)P)->PreSweeps = pre_sweeps;
				((ParBSRPrecondAMG*)P)->PostSweeps = post_sweeps;
				((ParBSRPrecondAMG*)P)->CoarseSweeps = coarse_sweeps;
				((ParBSRPrecondAMG*)P)->RelaxationType = relaxation_type;
				((ParBSRPrecondAMG*)P)->RelaxationFactor = relaxation_factor;
				((ParBSRPrecondAMG*)P)->PrintStats = print_stats;
			}
			else if (pre_type == 1)
			{
				P = new ParBSRPrecondILUAMG;

				((ParBSRPrecondILUAMG*)P)->MaxLevels = max_levels;
				((ParBSRPrecondILUAMG*)P)->CoarseSize = coarse_size;
				((ParBSRPrecondILUAMG*)P)->StrengthThreshold = strength_threshold;
				((ParBSRPrecondILUAMG*)P)->AggNumLevels = agg_num_levels;
				((ParBSRPrecondILUAMG*)P)->CoarsenType = coarsen_type;
				((ParBSRPrecondILUAMG*)P)->InterpMinElements = interp_min_elements;
				((ParBSRPrecondILUAMG*)P)->InterpMaxElements = interp_max_elements;
				((ParBSRPrecondILUAMG*)P)->TruncationFactor = truncation_factor;
				((ParBSRPrecondILUAMG*)P)->SparsificationThreshold = sparsification_threshold;
				((ParBSRPrecondILUAMG*)P)->CycleType = cycle_type;
				((ParBSRPrecondILUAMG*)P)->PreSweeps = pre_sweeps;
				((ParBSRPrecondILUAMG*)P)->PostSweeps = post_sweeps;
				((ParBSRPrecondILUAMG*)P)->CoarseSweeps = coarse_sweeps;
				((ParBSRPrecondILUAMG*)P)->PrintStats = print_stats;
			}
			else if (pre_type == 2)
				P = new ParBSRPrecondJacobi;
			else if (pre_type == 3)
				P = new ParBSRPrecondSOR(sor_type, sor_factor);
			else if (pre_type == 4)
				P = new ParBSRPrecondILU;

			setup_time = -MPI_Wtime();

			if (pre_type == 0)
				std::cout << "\nAMG Setup Phase:\n";
			else if (pre_type == 1)
				std::cout << "\nILUAMG Setup Phase:\n";
			else if (pre_type == 2)
				std::cout << "\nJacobi Setup Phase:\n";
			else if (pre_type == 3)
				std::cout << "\nSOR Setup Phase:\n";
			else if (pre_type == 4)
				std::cout << "\nILU Setup Phase:\n";
			
			std::cout << "----------------------------------------------------------------\n";

			((ParBSRPrecondBase*)P)->Setup(*(ParBSR*)A);

			setup_time += MPI_Wtime();

			std::cout << "Setup Time: " << setup_time << "\n";

			if (!rhs_file)
			{
				int N = A->OutSize() / 2;
				b.Resize(2 * N);
				double* bv = b.Loc.values;
				for (int i = 0; i < N; ++i)
				{
					bv[2 * i] = 1.0;
					bv[2 * i + 1] = 0.0;
				}
			}
			else
			{
				char line[1024];
				char banner[64];
				char mtx[64]; 
				char arr[64];
				char data_type[64];
				char storage_scheme[64];

				FILE* frhs = fopen(rhs_file, "r");

				fgets(line, 1024, frhs);
				sscanf(line, "%s %s %s %s %s", &banner, &mtx, &arr, &data_type, &storage_scheme);

				char* p;
				p = data_type;
				*p = toupper(*p);
				while (*(++p) != '\0')
					*p = tolower(*p);

				do
					fgets(line, 1024, frhs);
				while (line[0] == '%');

				if (strcmp(data_type, "Real") == 0)
				{
					int N, M;
					sscanf(line, "%d %d", &N, &M);

					double* bv = new double[2 * N];
					for (int i = 0; i < N; ++i)
					{
						fscanf(frhs, "%lg\n", &bv[2 * i]);
						bv[2 * i + 1] = 0.0;
					}
					
					fclose(frhs);
					
					b = ParVec(comm,2 * N, bv, 0);
				}
				else if (strcmp(data_type, "Complex") == 0)
				{
					int N, M;
					sscanf(line, "%d %d", &N, &M);

					double* bv = new double[2 * N];
					for (int i = 0; i < N; ++i)
						fscanf(frhs, "%lg %lg\n", &bv[2 * i], &bv[2 * i + 1]);
					
					fclose(frhs);
					
					b = ParVec(comm,2 * N, bv, 0);
				}
			}
		}

		x.Resize(A->InSize());
		x.Fill(0.0);

		solve_time = -MPI_Wtime();

		if (solver_type == 0)
			std::cout << "\nCG Solve Phase:\n";
		else if (solver_type == 1)
			std::cout << "\nGMRES Solve Phase:\n";
		else if (solver_type == 2)
			std::cout << "\nBiCGStab Solve Phase:\n";
		std::cout << "----------------------------------------------------------------\n";

		if (solver_type == 0)
			ParSolverCG(*A, *P, max_iters, tolerance, print_stats)(b, x, iter, relres);
		else if (solver_type == 1)
			ParSolverGMRES(*A, *P, gmres_restart, max_iters, tolerance, print_stats)(b, x, iter, relres);
		else if (solver_type == 2)
			ParSolverBCGS(*A, *P, max_iters, tolerance, print_stats)(b, x, iter, relres);

		solve_time += MPI_Wtime();

		std::cout << "Solve Time: " << solve_time << "\n";
		std::cout << "Time/Iteration: " << solve_time / iter << "\n";

		std::cout << "\n----------------------------------------------------------------\n";
		std::cout << "Total Time: " << setup_time + solve_time << "\n";

		if (out_file)
		{
			if (strcmp(data_type, "Real") == 0)
			{
				int N = x.Loc.size;
				double* xv = x.Loc.values;

				FILE* fout = fopen(out_file, "w");

				fprintf(fout, "%%%MatrixMarket matrix array real general\n");
				fprintf(fout, "%d %d\n", N, 1);

				for (int i = 0; i < N; ++i)
					fprintf(fout, "%20.16g\n", xv[i]);

				fclose(fout);
			}
			else if (strcmp(data_type, "Complex") == 0)
			{
				int N = x.Loc.size / 2;
				double* xv = x.Loc.values;

				FILE* fout = fopen(out_file, "w");

				fprintf(fout, "%%%MatrixMarket matrix array complex general\n");
				fprintf(fout, "%d %d\n", N, 1);

				for (int i = 0; i < N; ++i)
					fprintf(fout, "%20.16g %20.16g\n", xv[2 * i], xv[2 * i + 1]);

				fclose(fout);
			}
		}

		delete A;
		delete P;
	}

	MPI_Finalize();

	return 0;
}